#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""


import random



class Solution:
    def __init__(self):
        pass

    def topk(self, nums, k):
        if len(nums) < k:
            return nums

        self.quick_select(nums, k)

        return nums[:k]

    # 递归 quick_select: top(k) 放到nums[:k]
    def quick_select(self, nums, k):
        print(nums, k)
        assert nums
        pivot_idx = self.get_pivot_idx(nums)
        # pivot_idx = 0
        pivot_num = nums[pivot_idx]
        print("before ", pivot_idx, -1)
        nums[pivot_idx], nums[-1] = nums[-1], nums[pivot_idx]
        print("after ", nums)

        l = [1,2,3,4]
        print(l)
        l[0], l[-1] = l[-1], l[0]
        print(l)

        left_idx = 0

        for idx in range(len(nums)):
            # left --> right, max --> min
            if nums[idx] >= pivot_num:
                nums[left_idx], nums[idx] = nums[idx], nums[left_idx]
                left_idx += 1
            # else:
            #     pass

        print("left", left_idx, " nums:", nums, ", pivot:", pivot_num, ", pivot_idx", pivot_idx)
        print()
        if left_idx == k:
            return

        # 要选5个, 只获取到top 3
        if left_idx < k:
            self.quick_select(nums[left_idx:], k-left_idx)
        else:
            # 要选5个, 只获取到top 6
            # 从top6中获取top5
            self.quick_select(nums[:left_idx-1], k)


    def get_pivot_idx(self, nums):
        # 后面pivot可以灵活选取, 使用随机算法.
        # return len(nums) // 2
        return 0


def main():
    print("start main")
    s = Solution()
    # nums = [1, 2, 9, 8]
    nums = list(range(10))
    random.shuffle(nums)
    # nums[0], nums[-1] = nums[-1], nums[0]
    # nums = [8, 3, 2, 1, 4, 0, 7, 6, 5, 9]

    ret = s.topk(nums, 3)
    print(ret)
    print(nums)

if __name__ == "__main__":
    main()
